from mpl_toolkits import mplot3d
import numpy as np
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA

hyperbola = lambda x, a, b: np.sqrt(abs(b/a * (x**2 - a**2)))


def x_rotate_counter_clock(x, y, theta):
    xy = np.c_[x, y]
    theta_matrix = np.array([[np.cos(theta), np.sin(theta)],
                             [-np.sin(theta), np.cos(theta)]])
    xy2 = xy @ theta_matrix
    return xy2[:, 0], xy2[:, 1]


if '__main__' == __name__:
    fig = plt.figure(figsize=[16, 8])
    spr = 1
    spc = 2
    spn = 0
    spn += 1
    ax = fig.add_subplot(spr, spc, spn, projection='3d')
    n_dots = 16
    x_min = -6
    x_max = 6
    a = 2
    b = 3
    x = np.r_[np.linspace(x_min, -a, 1+n_dots), np.linspace(a, x_max, 1+n_dots)]

    # draw: hyperbola in xOy plane
    y = hyperbola(x, a, b)
    z = np.zeros_like(x)
    y_min = hyperbola(x_min, a, b)
    y_max = -y_min
    line_x = np.array([x_min, x_max, x_max, x_min, x_min])
    line_y = np.array([y_max, y_max, y_min, y_min, y_max])
    line_z = np.zeros_like(line_x)
    n_dots_border = len(line_x)
    ax.scatter3D(x, y, z, s=10, c='g')
    ax.scatter3D(x, -y, z, s=10, c='g')
    ax.plot3D(line_x, line_y, c='g')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    ax.grid()

    # draw: rotate 60 degree on axis y
    theta_by_y = np.pi / 3
    x2, z2 = x_rotate_counter_clock(x, z, theta_by_y)
    y2 = y
    ax.scatter3D(x2, y2, z2, s=10, c='y')
    ax.scatter3D(x2, -y2, z2, s=10, c='y')
    line_x2, line_z2 = x_rotate_counter_clock(line_x, line_z, theta_by_y)
    line_y2 = line_y
    ax.plot3D(line_x2, line_y2, line_z2, c='y')

    # draw: rotate 30 degree on axis z
    theta_by_z = np.pi / 6
    x3, y3 = x_rotate_counter_clock(x2, y2, theta_by_z)
    z3 = z2
    x3neg, y3neg = x_rotate_counter_clock(x2, -y2, theta_by_z)
    ax.scatter3D(x3, y3, z3, s=10, c='r')
    ax.scatter3D(x3neg, y3neg, z3, s=10, c='r')
    line_x3, line_y3 = x_rotate_counter_clock(line_x2, line_y2, theta_by_z)
    line_z3 = line_z2
    ax.plot3D(line_x3, line_y3, line_z3, c='r')

    # data to decompose
    xyz_pos = np.c_[x3, y3, z3]
    xyz_neg = np.c_[x3neg, y3neg, z3]
    xyz = np.r_[xyz_pos, xyz_neg]

    # decompose it
    dc = PCA(n_components=2)
    xy = dc.fit_transform(xyz)
    dcc = dc.components_
    print(dcc)
    print('explained var', dc.explained_variance_)
    print('explained var ratio', dc.explained_variance_ratio_)
    print('singular values', dc.singular_values_)
    cmap = plt.cm.get_cmap('rainbow', len(dcc))

    # draw the 特征向量 of the decomposition
    for i, dcc_vec in enumerate(dcc):
        dcc_vec *= 8
        ax.plot3D([0, dcc_vec[0]], [0, dcc_vec[1]], [0, dcc_vec[2]], '--', c=cmap(i))

    # draw: the decomposed data
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.scatter(xy[:, 0], xy[:, 1])
    plt.grid()
    plt.xlim([-8, 8])
    plt.ylim([-8, 8])

    plt.show()
